import json
import os
import shapely.geometry as geom
import numpy as np


obj_list = ["WHITE_LABEL"]

obj_dict = {
    "WHITE_LABEL": 0,
}


def is_valid(point):
    """
    Check if the point is valid
    """
    if point[0] < 0 or point[0] > 1 or point[1] < 0 or point[1] > 1:
        return 0
    return 1


if __name__ == "__main__":
    obj_dir = r"/home/ubuntu/cxf/yolo_train/datasets/barcode_yolo_obb_label/"
    name_list = os.listdir(obj_dir)
    for name1 in name_list:
        if name1.split(".")[1] == "json":
            json_path = os.path.join(obj_dir, name1.split(".")[0] + ".json")
            txt_path = os.path.join(obj_dir, name1.split(".")[0] + ".txt")

            row_list = []
            with open(json_path) as f:
                dict1 = json.load(f)
                imageWidth = dict1["imageWidth"]
                imageHeight = dict1["imageHeight"]
                img_name = dict1["imagePath"]
                shapes = dict1["shapes"]
                rectangle = np.array([(0, 0), (1, 0), (1, 1), (0, 1)])
                rectangle_geom = geom.Polygon([(0, 0), (1, 0), (1, 1), (0, 1)])
                for shape in shapes:
                    if shape["label"] in obj_list:
                        x1 = shape["points"][0][0] / imageWidth
                        y1 = shape["points"][0][1] / imageHeight
                        x2 = shape["points"][1][0] / imageWidth
                        y2 = shape["points"][1][1] / imageHeight
                        x3 = shape["points"][2][0] / imageWidth
                        y3 = shape["points"][2][1] / imageHeight
                        x4 = shape["points"][3][0] / imageWidth
                        y4 = shape["points"][3][1] / imageHeight
                        polygon = geom.Polygon(
                            [(x1, y1), (x2, y2), (x3, y3), (x4, y4)]
                        )
                        # 求交集
                        intersection = polygon.intersection(rectangle_geom)
                        # 获取交集的 bounding box
                        minx, miny, maxx, maxy = intersection.bounds
                        bounding_box = [(minx, miny), (maxx, maxy)]
                        # 若没有交集，则 bouding box 中出现 nan
                        if not ("nan" in str(bounding_box)):
                            print("Bounding Box:", bounding_box)
                            center_x = (minx + maxx) / 2
                            center_y = (miny + maxy) / 2
                            width = maxx - minx
                            height = maxy - miny
                            row1 = [
                                0,  # 占位的类别，后续修改
                                center_x,
                                center_y,
                                width,
                                height,
                            ]
                            for i in range(len(row1)):
                                if row1[i] > 1:
                                    row1[i] = 1
                                if row1[i] < 0:
                                    row1[i] = 0
                            # 写入 txt 文件
                            row1 = list(map("{:.6f}".format, row1))
                            row1[0] = obj_dict[shape["label"]]
                            row_list.append(row1)

            with open(txt_path, "w", newline="") as f:
                for row in row_list:
                    for data in row:
                        f.write(str(data))
                        f.write(" ")
                    f.write("\r\n")
